import os
import datetime

from common.geometry.geom import Geom
from common.geoimage.gdal_sys import (
    mosaic_ds,
    translate_ds,
    warp_ds,
    ADDITIONAL_GEOTIFF_LERC_ZSTD_PROFILE,
)
from common.geoimage.raster_dataset import (
    RasterDataset,
    stack_raster_datasets,
)

import catalog
from catalog.utils import interpret_protocol
from catalog.models.sentinel2_l2a import Sentinel2L2AScene


BASE_PATH = "/media/workspace/zhouhang/sample_data"


def test_mosaic_multiple_ds_to_file():
    fpath_loc = os.path.join(BASE_PATH, "china_geojson_eng/xian4/fu4yu4xian4.geojson")
    aoi = Geom.from_geojson_fpath(fpath_loc)

    print("query to get the items")
    items = Sentinel2L2AScene.query_many_items(
        spatial=aoi,
        temporal=[datetime.datetime(2021, 7, 10), datetime.datetime(2021, 7, 11)],
    )
    items = list(items)
    print(items)

    print("mosaic ds")
    lst_fpaths = [
        interpret_protocol(item.info["assets"]["B08"]["href"]) for item in items
    ]
    print(lst_fpaths)

    mosaic_ds(
        lst_fpaths,
        resample_alg="average",
        src_nodata=0,
        dest_res=100,
        fpath_dest="test1_mosaiced_b08_100m_lerczstd.tif",
        options=ADDITIONAL_GEOTIFF_LERC_ZSTD_PROFILE,
    )
    print("mosaic raster dataset")


def test_stack_multiple_raster_dataset():
    cid = catalog.meta.list_collections(filters=["SENTINEL2"])[0]
    print(cid)

    print("query to get the items")
    items = catalog.meta.query_many_items(
        collection_id=cid,
        temporal=[
            datetime.datetime(2021, 7, 10),
            datetime.datetime(2021, 8, 10, 23, 59, 59),
        ],
        properties={"mgrs_tile": "51UXP"},
    )
    items = list(items)
    print(items)

    lst_rds = []
    for item in items[0:4]:
        rds = catalog.raster.retrieve_by_item(item, asset_names=["B04", "B03", "B02"])
        lst_rds.append(rds)
    print("all retrieved")
    print(lst_rds)

    print("stack ds")
    result_rds = stack_raster_datasets(lst_rds)
    print(result_rds.meta)
    print(result_rds.data.shape)


def test_translate_to_file():
    fpath_loc = os.path.join(BASE_PATH, "china_geojson_eng/xian4/fu4yu4xian4.geojson")
    aoi = Geom.from_geojson_fpath(fpath_loc)

    print("query to get the items")
    items = Sentinel2L2AScene.query_many_items(
        spatial=aoi,
        temporal=[datetime.datetime(2021, 7, 10), datetime.datetime(2021, 7, 11)],
    )
    item = list(items)[0]

    fpath_input = interpret_protocol(item.info["assets"]["B04"]["href"])
    print(fpath_input)
    translate_ds(fpath_input, aoi=aoi, fpath_dest="./test3_translate_to_file.tif")


def test_warp_to_file():
    fpath_loc = os.path.join(BASE_PATH, "china_geojson_eng/xian4/fu4yu4xian4.geojson")
    aoi = Geom.from_geojson_fpath(fpath_loc)

    print("query to get the items")
    items = Sentinel2L2AScene.query_many_items(
        spatial=aoi,
        temporal=[datetime.datetime(2021, 7, 10), datetime.datetime(2021, 7, 11)],
    )
    item = list(items)[0]

    fpath_input = interpret_protocol(item.info["assets"]["B04"]["href"])
    print(fpath_input)

    from osgeo import gdal

    ds = gdal.Open(fpath_input)
    bd = ds.GetRasterBand(1)
    bd.SetNoDataValue(10)
    del bd

    warp_ds(
        # fpath_input,
        ds,
        dest_srs="EPSG:4326",
        dest_extent=aoi.bounds,
        # dest_nodata=0,
        fpath_dest="./test4_warp_to_file.tif",
    )


def test_warp_with_multiple_dest_nodata():
    fpath_input = (
        "/NAS5/agri/chendu/crop_sys/decloud_rasters/days_16_False_year_split/T10/T10SDH/2019/"
        "decloud_10m_10SDH_20190117_20190202.jp2"
    )
    meta = RasterDataset.read_meta(fpath_input)
    warp_ds(
        fpath_input,
        dest_srs=meta.prj_epsg,
        dest_nodata=[0, 1, 2, 3, 4, 5, 6],
        fpath_dest="./test5_different_dest_nodata_values.tif",
    )


def main(option):
    if option == 1:
        test_mosaic_multiple_ds_to_file()
    elif option == 2:
        test_stack_multiple_raster_dataset()
    elif option == 3:
        test_translate_to_file()
    elif option == 4:
        test_warp_to_file()
    elif option == 5:
        test_warp_with_multiple_dest_nodata()


if __name__ == "__main__":
    import fire

    fire.Fire(main)
